import matplotlib.pyplot as plt
import numpy as np


def mha_visualization(weights, words=None, save_name=None):
    """
    Attention Visualization
    Args:
        weights: [l, t]
    """
    
    assert weights.ndim == 3

    nhead = weights.shape[0]
    assert nhead > 1
    
    plt.figure(figsize=[15, 20])
    for i in range(1, nhead+1):
        ax = plt.subplot(int('%d1%d' % (nhead+1, i)))
        ax.matshow(weights[i-1])
        ax.set_title('The Attention Weights of the No.%d Head' % i)
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')

        # Set up axes
        if words is not None:
            ax.set_yticklabels([''] + words.split() + ['<EOS>'])

    # plot the mean of weight
    ax = plt.subplot(int('%d1%d' % (nhead+1, nhead+1))) 
    ax.matshow(np.mean(weights, axis=0))
    ax.set_title('The Mean of Attention Weights')
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    if words is not None:
        ax.set_yticklabels([''] + words.split() + ['<EOS>'])

    if save_name is not None:
        plt.savefig(save_name)
        plt.close()
    else:
        plt.show()


def mha_visualization_mean(weights, words=None, save_name=None):
    """
    Attention Visualization
    Args:
        weights: [l, t]
    """
    
    assert weights.ndim == 3

    nhead = weights.shape[0]
    assert nhead > 1
    
    if weights.shape[1] == weights.shape[2]:
        plt.figure(figsize=[5, 5])
    else:
        plt.figure(figsize=[15, 4])

    # plot the mean of weight
    ax = plt.subplot(111)
    ax.matshow(np.mean(weights, axis=0))
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    if words is not None:
        ax.set_yticklabels([''] + words.split() + ['<EOS>'])

    if weights.shape[1] == weights.shape[2]:
        ax.set_xlabel('#Label')
        ax.set_title('The Visualization of Self-Attention Weights')
    else:    
        ax.set_xlabel('#Frame')
        ax.set_title('The Visualization of EncDec-Attention Weights')

    ax.set_ylabel('#Label')

    if save_name is not None:
        plt.savefig(save_name)
        plt.close()
    else:
        plt.show()